import argparse
import os

from rst import RSTNode, RSTTree
from relation_set import rstdt_mapping18

def depth_node(node): 
    depth = 0
    curent_node = node
    while curent_node.parent_node is not None:
        depth += 1
        curent_node = curent_node.parent_node
    node.depth = depth
    for child_node in node.child_nodes:
        depth_node(child_node)
    return

def braching_nodes(rst_tree):
    branching(rst_tree, rst_tree.root_node)
    depth_node(rst_tree.root_node)


def branching(tree, node):
    if len(node.child_nodes) > 2:
        tree.branching_count += 1

        n_count = 0
        for idx, child_node in enumerate(node.child_nodes): 
            if child_node.nuclearity == 'Nucleus':
                n_count += 1

        if n_count == len(node.child_nodes):
            multinuc_flag = True
            ref_relation = node.child_nodes[0].relation
            for child_node in node.child_nodes:
                assert child_node.relation == ref_relation
        else:
            multinuc_flag = False

        ## branching 
        while len(node.child_nodes) > 2:
            n_index = -1
            n_count = 0
            for idx, child_node in enumerate(node.child_nodes): 
                if child_node.nuclearity == 'Nucleus':
                    n_count += 1
                    n_index = idx
            
            new_node = RSTNode()
            rst_tree.all_nodes.append(new_node)
            new_node.nuclearity = 'Nucleus'

            if multinuc_flag:
                new_node.relation = ref_relation
            else:
                new_node.relation = 'span'
                assert n_count == 1

            if n_index == len(node.child_nodes) - 1:
                left_node =  node.child_nodes[n_index - 1]
                right_node = node.child_nodes[n_index]
            else:
                left_node = node.child_nodes[n_index]
                right_node =  node.child_nodes[n_index + 1]

            new_node.edu_start = left_node.edu_start 
            new_node.edu_end = right_node.edu_end
            new_node.child_nodes.append(left_node)
            new_node.child_nodes.append(right_node)
            left_node.parent_node = new_node
            right_node.parent_node = new_node
            node.child_nodes.remove(left_node)
            node.child_nodes.remove(right_node)

            new_node.parent_node = node
            for position, child_node in enumerate(node.child_nodes):
                if child_node.edu_start > new_node.edu_start:
                    break
            node.child_nodes.insert(position + 1, new_node)

    for child_node in node.child_nodes: 
        branching(tree, child_node)
    return

def info(line, depth):
    info_list = line.split(" ")
    node_type = info_list[2][1:]

    rst_node = RSTNode()
    rst_node.depth = depth

    if node_type == "leaf":

        nuc_str = info_list[1]
        edu_id = info_list[3][:-1]
        rel_str = info_list[5][:-1]

        assert info_list[4] == '(rel2par'
        assert info_list[6] == '(text'
        assert edu_id.isnumeric()

        text_info = info_list[7:-1]
        text_info[0] = text_info[0][2:]
        text_info[-1] = text_info[-1][:-3]

        rst_node.edu_start = int(edu_id)
        rst_node.edu_end = int(edu_id)
        rst_node.nuclearity = nuc_str
        rst_node.relation = rel_str
        rst_node.text = text_info

    elif node_type == "span":
        nuc_str = info_list[1]
        edu_start = info_list[3]
        edu_end = info_list[4][:-1]

        if nuc_str != 'Root':
            assert info_list[5] == '(rel2par'
            rel_str = info_list[6][:-1]

            rst_node.edu_start = int(edu_start)
            rst_node.edu_end = int(edu_end)
            rst_node.nuclearity = nuc_str
            rst_node.relation = rel_str

        else: ### elif nuc_str == 'Root'
            rst_node.edu_start = int(edu_start)
            rst_node.edu_end = int(edu_end)
            rst_node.nuclearity = nuc_str
    return rst_node

def get_depth(line):
    depth = 0
    for ch in line:
        if ch == ' ':
            depth += 1
        else:
            break
    depth = depth // 2
    return depth

def read_original_dis(file_path):
    with open(file_path, mode='r', encoding='utf8') as inf:
        rst_tree = RSTTree()

        edu_stack = []
        for idx, line in enumerate(inf.readlines()):
            clean_line = line.strip()
            depth = get_depth(line)

            if clean_line[0] != ")":
                rst_node = info(clean_line, depth)
                edu_stack.append(rst_node)

                if rst_node.is_leaf():
                    rst_tree.leaf_nodes.append(rst_node)
                rst_tree.all_nodes.append(rst_node)
            else:
                child_nodes = []
                while edu_stack[-1].depth != depth:
                    child_nodes.append(edu_stack[-1])
                    edu_stack.pop()
                child_nodes.reverse()

                edu_stack[-1].child_nodes = child_nodes

                for child_node in child_nodes:
                    child_node.parent_node = edu_stack[-1]
                
                assert child_nodes[0].edu_start == edu_stack[-1].edu_start
                assert child_nodes[-1].edu_end == edu_stack[-1].edu_end
    assert len(edu_stack) == 1
    rst_tree.root_node = edu_stack[0]
    return  rst_tree

def check(rst_tree):
    #branching(root_node)
    check_child_nodes(rst_tree.root_node)
    leaf_nodes = []
    gather_leaf_nodes(rst_tree.root_node, leaf_nodes)
    assert len(leaf_nodes) == len(rst_tree.leaf_nodes)
    assert len(rst_tree.leaf_nodes) * 2 - 1 == len(rst_tree.all_nodes)
    return

def check_child_nodes(node):
    if len(node.child_nodes) > 0:
        assert node.child_nodes[0].edu_start == node.edu_start
        assert node.child_nodes[-1].edu_end == node.edu_end
        assert len(node.child_nodes) == 2
        assert node.child_nodes[0].edu_end + 1 == node.child_nodes[1].edu_start
        nn_str = node.child_nodes[0].nuclearity + "#" + node.child_nodes[1].nuclearity
        assert nn_str in {'Nucleus#Nucleus', 'Nucleus#Satellite', 'Satellite#Nucleus'}

    for child_node in node.child_nodes: 
        check_child_nodes(child_node)

def gather_relations(rst_tree, rel_set):
    get_relation(rst_tree.root_node, rel_set)
    return

def get_relation(node, rel_set):
    if node.relation is not None:
        rel_set.add(node.relation)

    for child_node in node.child_nodes:
        get_relation(child_node, rel_set)
    return

def gather_leaf_nodes(node, leaf_nodes):
    if len(node.child_nodes) == 0:
        leaf_nodes.append(node)

    for child_node in node.child_nodes: 
        gather_leaf_nodes(child_node, leaf_nodes)

def mapping(tree):
    relation_mapping(tree.root_node)
    return

def relation_mapping(node):
    if node.nuclearity != 'Root' and node.relation != "span":
        mapped_relation = rstdt_mapping18.get(node.relation.lower())
        assert mapped_relation is not None
        node.relation = mapped_relation
    
    for child_node in node.child_nodes:
        relation_mapping(child_node)
    return

if __name__ == '__main__':
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--input_dir', default='train')
    argparser.add_argument('--output_dir', default='train_r')
    argparser.add_argument('--map', action='store_true')
    args, extra_args = argparser.parse_known_args()

    if not os.path.exists(args.output_dir):
        os.mkdir(args.output_dir)

    overall_branching_count = 0
    overall_non_leaf_node_count = 0
    overall_leaf_node_count = 0

    file_count = 0
    rel_set = set()
    for file_name in os.listdir(args.input_dir):
        if file_name[-3:] != 'dis': continue
        file_path = os.path.join(args.input_dir, file_name)
        print("file: ", file_path)
        rst_tree = read_original_dis(file_path)        

        braching_nodes(rst_tree)
        file_count += 1
        overall_branching_count += rst_tree.branching_count 
        overall_non_leaf_node_count += (len(rst_tree.all_nodes) - len(rst_tree.leaf_nodes))
        overall_leaf_node_count += len(rst_tree.leaf_nodes)
        if args.map:
            mapping(rst_tree)

        gather_relations(rst_tree, rel_set)
        check(rst_tree)
        output_path = os.path.join(args.output_dir, file_name)
        rst_tree.save(output_path)

    print("file: ", file_count) 
    # print("branching nodes:", overall_branching_count)
    print("non-leaf nodes:", overall_non_leaf_node_count)
    print("leaf nodes:", overall_leaf_node_count)

    print("relation set:", rel_set)
    print("relation num: ", len(rel_set))